package com.sluggard.ai.manage;

import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import org.springframework.ai.mcp.AsyncMcpToolCallbackProvider;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpStdioClientProperties;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 类描述：
 *
 * @author lizheng
 * @created 2025/4/27 15:52
 */
public class DynamicMcpManager {

    private static final ConcurrentMap<String, McpAsyncClient> mcpAsyncClientMap = new ConcurrentHashMap<>();

    public static AsyncMcpToolCallbackProvider asyncMcpToolCallbackProvider(List<String> serverNames) {
        if(Objects.isNull(serverNames) || serverNames.isEmpty()){
            return null;
        }
        return new AsyncMcpToolCallbackProvider(serverNames.stream()
                .map(mcpAsyncClientMap::get)
                .toList());
    }

    public static List<McpAsyncClient> getMcpAsyncClients(List<String> serverNames) {
        return serverNames.stream()
                .map(mcpAsyncClientMap::get)
                .toList();
    }

    public static McpAsyncClient getMcpAsyncClient(String serverName) {
        return mcpAsyncClientMap.get(serverName);
    }

    public static List<McpAsyncClient> getMcpAsyncClients() {
        return new ArrayList<>(mcpAsyncClientMap.values());
    }

    public static McpAsyncClient buildMcpAsyncClient(ServerParameters serverParameters) {
        var mcpClient = McpClient.async(new StdioClientTransport(serverParameters))
                .requestTimeout(Duration.ofSeconds(10)).build();
        mcpClient.initialize().block();
        return mcpClient;
    }

    public static void updateServers(Map<String, McpStdioClientProperties.Parameters> mcpServers) {
        mcpServers.forEach((serverName, serverParameters) -> {
            mcpAsyncClientMap.putIfAbsent(serverName, buildMcpAsyncClient(serverParameters.toServerParameters()));
        });
    }
}
